#! /usr/bin/env

import torch
from torch import nn
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

training_data = datasets.FashionMNIST(root="data", train=True, download=True, transform=ToTensor())
test_data = datasets.FashionMNIST(root="data", train=False, download=True, transform=ToTensor())

# print(training_data.__len__)
# print(test_data.__len__)
# print(training_data.__getitem__(0))
# print(test_data.__getitem__(0))

#可视化训练数据
labels_map = {
    0: "T-Shirt",
    1: "Trouser",    
    2: "Pullover",    
    3: "Dress",    
    4: "Coat",    
    5: "Sandal",    
    6: "Shirt",    
    7: "Sneaker",    
    8: "Bag",    
    9: "Ankle Boot",
}
figure = plt.figure(figsize=(8,8))
cols,rows = 3,2
for i in range(1, cols * rows + 1):
    sample_index = torch.randint(len(training_data), size=(1,)).item()
    sample_index = i
    # print(sample_index)
    img,lable = training_data[sample_index]
    print(img.shape)
    # print(labels_map[lable])
    figure.add_subplot(rows*2,cols,i)
    plt.title(labels_map[lable])
    plt.axis("off")
    plt.imshow(img.squeeze(),cmap="gray")

for i in range(1, cols * rows + 1):
    sample_index = torch.randint(len(test_data), size=(1,)).item()
    sample_index = i
    # print(sample_index)
    img,lable = test_data[sample_index]
    # print(img.shape)
    # print(labels_map[lable])
    figure.add_subplot(rows*2,cols,i + rows * cols)
    plt.title(labels_map[lable])
    plt.axis("off")
    plt.imshow(img.squeeze(),cmap="Blues")

plt.show()

#将数据集疯转到dataloader中，方便在训练过程中迭代处理数据
train_dataloader = DataLoader(training_data, batch_size=3, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=3, shuffle=True)

for epoch_num in range(2):
    for batch_idx, (data, labels) in enumerate(train_dataloader):
        print(f"Batch {batch_idx + 1}:")
        print(f"Data: {data}")
        print(f"Labels: {labels}")